import mindspore as ms
from mindspore import nn


ms.set_context(device_target = 'CPU', mode = ms.PYNATIVE_MODE)


class SimpleNet(nn.Cell):
    def __init__(self,in_channel = 3,clses = 10, h = 32, w = 32,auto_prefix=True, flags=None):
        super().__init__(auto_prefix, flags)

        self.w = w
        self.h = h

        self.act = nn.ReLU()
        self.flatten = nn.Flatten()

        self.conv1 = nn.Conv2d(in_channel, 6, 5)
        self.pool1 = nn.MaxPool2d(2,2)


        self.conv2 = nn.Conv2d(6, 16, 5)
        self.pool2 = nn.MaxPool2d(2,2)

        # self.conv3 = nn.Conv2d(8,4,3)
        # self.pool3 = nn.AvgPool2d(2, 2)


        self.fc = nn.Dense(self.get_feature(4) * 16, 64)

        self.fc1 = nn.Dense(64, 64)

        self.out = nn.Dense(64, clses)


    def get_feature(self, ratio):
        return  int(self.h / ratio) * int(self.w / ratio)

    def construct(self, *inputs):
        x = self.act(self.pool1(self.conv1(*inputs)))

        x = self.act(self.pool2(self.conv2(x)))

        # x = self.act(self.pool3(self.conv3(x)))

        x = self.act(self.fc(self.flatten(x)))

        x = self.fc1(x)
        x = self.out(x)

        return x



def generate_model(net:nn.Cell, sets, epoch = 1000):
    loss_fn = nn.SoftmaxCrossEntropyWithLogits(True,"mean")
    optimizer = nn.optim.SGD(net.trainable_params(), learning_rate=1e-2, momentum=0.9)

    model = ms.Model(net, loss_fn, optimizer, metrics={"recall":nn.Recall(), "precision":nn.Precision()})

    model.train(epoch, sets['train'], callbacks=ms.LossMonitor(40))

    out = model.eval(sets['test'])

    print(out)


class DQN(nn.Cell):
    def __init__(self,in_dim,out_dim = 4,auto_prefix=True, flags=None):
        super().__init__(auto_prefix, flags)
        self.in_dim = in_dim
        self.out_dim = out_dim

        self.fc = nn.SequentialCell(
            nn.Dense(in_dim, 8),
            nn.ReLU(),
            nn.BatchNorm1d(8),
            nn.Dense(8, out_dim)
        )

    def construct(self, *inputs):
        x = self.fc(*inputs)
        return x

if __name__ == "__main__":
    pass
